#!/usr/bin/env python3
print("Loading modules...")
import argparse
import pandas as pd
from DeepPurpose.DTI import model_pretrained, virtual_screening


def parse_arguments():
    parser = argparse.ArgumentParser(
        description="Virtual Screening using a pre-trained model."
    )
    parser.add_argument(
        "-M",
        "--model",
        required=True,
        help="Path to the folder containing the pre-trained model",
    )
    parser.add_argument(
        "-R",
        "--receptor",
        required=True,
        help="FASTA file containing the protein amino acid sequence",
    )
    parser.add_argument(
        "-L",
        "--ligand",
        required=True,
        help="SMI file containing ligand SMILES sequences",
    )
    parser.add_argument("-O", "--output", required=True, help="Output CSV file name")
    return parser.parse_args()


def read_fasta(file_path):
    with open(file_path, "r") as file:
        lines = file.readlines()
        receptor_name = lines[0].strip()[1:]  # Remove '>' and newline character
        receptor_sequence = "".join(lines[1:]).replace(
            "\n", ""
        )  # Join all lines and remove newlines
    return [receptor_sequence], receptor_name


def read_smi(file_path):
    # Try to read with header
    try:
        df = pd.read_csv(file_path, sep=" ", header=0)
        if list(df.columns) == ["smiles", "zinc_id"]:
            ligands = df["smiles"].tolist()
            ligand_names = df["zinc_id"].tolist()
            return ligands, ligand_names
    except Exception as e:
        pass

    # If reading with header fails, read without header
    df = pd.read_csv(file_path, sep=" ", header=None)
    if len(df.columns) != 2:
        raise ValueError(
            "SMI file must have exactly two columns (SMILES and Ligand Name)."
        )
    ligands = df.iloc[:, 0].tolist()
    ligand_names = df.iloc[:, 1].tolist()
    return ligands, ligand_names


def main():
    args = parse_arguments()

    # Load the pre-trained model
    net = model_pretrained(args.model)

    # Read receptor sequence
    receptor, receptor_name = read_fasta(args.receptor)

    # Read ligand sequences and names
    ligands, ligand_names = read_smi(args.ligand)

    # Perform virtual screening
    res = virtual_screening(ligands, receptor, net, ligand_names, receptor_name)

    # Prepare output data
    output_data = {"SMILES": ligands, "Ligand_Name": ligand_names, "Score": res}

    # Write results to CSV
    output_df = pd.DataFrame(output_data)
    output_df.to_csv(args.output, index=False)


if __name__ == "__main__":
    main()
